// This is the main DLL file.

#include "gi_sample.h"

#include "D3Dx11.h"
#include "D3DX10math.h"

#include "RBApplication/RBApplication.h"
#include "RBMain/RBMaterial.h"
#include "RBMain/RBGeometry.h"
#include "RBMain/RBRenderTargetGroup.h"
#include "RBMain/RBRenderPassFullScreen.h"
#include "RBMain/RBRenderPass.h"
#include "RBMain/RBRenderPassDebug.h"
#include "RBRenderProcess/RBEffectFactory.h"
#include "sh_octree.h"
#include "RBMain/sh_math.h"
#include "RBD3D11Renderer/RBD3D11Renderer.h"
#include "RBRenderProcess/RBRenderTargetView.h"

struct rb_private_members
{
    RBGeometryHandle m_debug_sphere_template;
    std::vector<RBGeometryHandle> m_debug_spheres;
    RBEntityHandle m_debug_sphere_entity;
    sh_octree octree;
    ID3D11Query* m_query;
};

void read_rgb_hdr(const void* pixel_value, float* r, float*g, float* b)
{
    D3DXVECTOR4_16F* data = (D3DXVECTOR4_16F*)pixel_value;
    *r = data->x * 0.17f;
    *g = data->y * 0.17f;
    *b = data->z * 0.17f;
}

void read_rgb_ldr(const void* pixel_value, float* r, float*g, float* b)
{
    DWORD* data = (DWORD*)pixel_value;
    *r = ((*data & 0x000000FF) >>  0) / 255.0f;
    *g = ((*data & 0x0000FF00) >>  8) / 255.0f;
    *b = ((*data & 0x00FF0000) >> 16) / 255.0f;
}

void write_rgb_ldr(void* pixel_value, float r, float g, float b)
{
    DWORD* data = (DWORD*)pixel_value;
    *data =  static_cast<uint8>(RBMath::Clamp(r*255.0f, 0.0f, 255.0f)) | 
        (static_cast<uint8>(RBMath::Clamp(g*255.0f, 0.0f, 255.0f)) << 8) | 
        (static_cast<uint8>(RBMath::Clamp(b*255.0f, 0.0f, 255.0f)) << 16);
}

void projectCubeMap(rb_locked_data data[], size_t res, float* sh_r, float* sh_g, float* sh_b, int order)
{
    void* datas[6];
    for (int32 ui=0; ui < 6; ++ui)
    {
        datas[ui] = data[ui].Data;
    }
    RB::sh_math::project_cube_map(sh_r, sh_g, sh_b, datas,  data[0].Pitch, res, sizeof(DWORD), order, read_rgb_ldr);
    //RB::sh_math::project_cube_map(sh_r, sh_g, sh_b, datas,  data[0].Pitch, res, sizeof(D3DXVECTOR4_16F), order, read_rgb_hdr);
}

// needs order^2 coeffs
void SH_Encode(RBTexture* tex, float* sh_r, float* sh_g, float* sh_b, int order)
{
    // only up to order 3 is supported right now

    rb_locked_data data[6];
    uint32 mip_count = tex->GetMipCount();

    for (int i=0; i<6; ++i)
    {
        UINT sub_res = D3D10CalcSubresource(0, i, mip_count);
        RBRenderer::GetSingleton()->Lock(tex, D3D11_MAP_READ, &data[i], sub_res);
    }

    projectCubeMap(data, tex->GetHeight(), sh_r, sh_g, sh_b, order);

    for (int i=0; i<6; ++i)
    {
        UINT sub_res = D3D10CalcSubresource(0, i, mip_count);
        RBRenderer::GetSingleton()->Unlock(tex, sub_res);
    }
}

void SH_Decode(RBTexture* tex, RBTexture* scratch_tex, RBTexture* src_tex, const float* sh_r, const float* sh_g, const float* sh_b, int order, bool use_sh)
{
    rb_locked_data data;
    uint32 tex_res = tex->GetHeight();

    uint32 src_face;

    for (int face=0; face<6; ++face)
    {      
        src_face = face;
        if (use_sh)
        {
            RBRenderer::GetSingleton()->Lock(scratch_tex, D3D11_MAP_WRITE_DISCARD, &data);
            RB::sh_math::unproject_cube_map_face(sh_r, sh_g, sh_b, data.Data, face, data.Pitch, tex_res, sizeof(DWORD), order, write_rgb_ldr);
            RBRenderer::GetSingleton()->Unlock(scratch_tex);
            src_tex = scratch_tex;
            src_face = 0;
        }

        UINT src_sub_res = D3D10CalcSubresource(0, src_face, src_tex->GetMipCount());
        UINT dst_sub_res = D3D10CalcSubresource(0, face, tex->GetMipCount());


        RBD3D11Renderer::GetSingleton()->GetDeviceContext()->CopySubresourceRegion((ID3D11Resource*)tex->GetNative(), dst_sub_res, 0, 0, 0, (ID3D11Resource*)src_tex->GetNative(), src_sub_res, NULL);
    }

}

volatile bool test_sh = true;
float3 g_dirs[] = 
{   
    float3(1,0,0), float3(-1,0,0),
    float3(0,1,0), float3(0,-1,0),
    float3(0,0,1), float3(0,0,-1)
};

float3 g_ups[] = 
{   
    float3(0,1,0), float3(0,1,0),
    float3(0,0,-1), float3(0,0,1),
    float3(0,1,0), float3(0,1,0)
};

volatile int cam_idx=0;

volatile bool dump_profiler = false;

UINT64 render_cube(float3 pos, RBCamera* camera, RBRenderTargetGroup* target, ID3D11Query* query, uint32 mask=0xFFFFFFFF)
{
    UINT64 pixel_passed = 0;
    camera->WorldXForm().Translate = pos;
    for (uint32 ui=0; ui < 6; ++ui)
    {
        if (((1<<ui) & mask) == 0) continue;

        camera->LookAt(g_dirs[ui], g_ups[ui]);
        // Go left handed because we handle with a reflection
        {
            float* col_x = camera->WorldXForm().Rotate.GetCol(RBCamera::CD_RIGHT);
            col_x[0] = -col_x[0];
            col_x[1] = -col_x[1];
            col_x[2] = -col_x[2];
        }

        RBGeometryArray geoms;
        {
            PROFILE_BLOCK(RBApplication_Visibility);
            rb_render_pass::collect_geometries(&geoms, camera, RBApplication::GetSingleton()->get_scene());
        }

        {
            PROFILE_BLOCK(RBApplication_Render);
            rb_render_pass::open_render_target_group(RBRenderer::GetSingleton(), target, RBRenderer::CF_ALL, RBColorA(0,1,1,1));
            RBD3D11Renderer::GetSingleton()->GetDeviceContext()->Begin(query);
            rb_render_pass::render_geometries(RBRenderer::GetSingleton(), geoms, camera);
            RBD3D11Renderer::GetSingleton()->GetDeviceContext()->End(query);
        }

        RBD3D11Renderer::GetSingleton()->Flush();

        UINT64 queryData; // This data type is different depending on the query type
        while( S_OK != RBD3D11Renderer::GetSingleton()->GetDeviceContext()->GetData(query, &queryData, sizeof(UINT64), 0) )
        {
            ::Sleep(1);
        }

        pixel_passed += queryData;
    }

    return pixel_passed;

}


static const float3 factors[] = {
    float3(-0.5f, -0.5f, -0.5f),
    float3(-0.5f, -0.5f,  0.5f),
    float3(-0.5f,  0.5f, -0.5f),
    float3(-0.5f,  0.5f,  0.5f),
    float3( 0.5f, -0.5f, -0.5f),
    float3( 0.5f, -0.5f,  0.5f),
    float3( 0.5f,  0.5f, -0.5f),
    float3( 0.5f,  0.5f,  0.5f)
};

static const RBColor colors[] = {
    RBColor(0.0f, 0.0f, 0.0f),
    RBColor(1.0f, 0.0f, 0.0f),
    RBColor(0.0f, 1.0f, 0.0f),
    RBColor(0.0f, 0.0f, 1.0f),
    RBColor(1.0f, 1.0f, 0.0f),
    RBColor(0.0f, 1.0f, 1.0f),
    RBColor(1.0f, 0.0f, 1.0f),
    RBColor(1.0f, 1.0f, 1.0f)
};



volatile int draw_aab_count = 0;

void draw_octree(const sh_octree* octree, float3 center, float3 extents, uint16 cur_node = 0, uint16 depth=0)
{
    if (octree->m_nodes[cur_node].is_leaf())
    {
        {
            RBRenderPassDebug::add_aabb(center, extents, colors[depth]);
        }
    }
    else
    {
        stb_uint16 child_idx = octree->m_nodes[cur_node].get_child();
        for (int i=0; i<8; ++i)
        {
            draw_octree(octree, center + extents * factors[i], extents * 0.5f, child_idx+i, depth+1);
        }
    }
}

void build_sh_octree(sh_octree* octree, RBRenderTargetGroup* target, float3 center, float extents, RBCamera* camera, ID3D11Query* query, uint16 cur_node = 0, int depth=0)
{
    if (depth == 5)
    {
        return;
    }

    static const float ratio = 8000.0f;
    static const float min_near = 0.0001f;
    camera->SetViewFrustum(-1.0f, 1.0f, 1.0f, -1.0f, RBMath::Max(min_near, extents/ratio), extents, false);

    UINT64 pixel_count = render_cube(center, camera, target, query);

    // query has issued: how many pixel were written ?
    if (pixel_count > 0)
    {
        if (octree->subdivide(cur_node))
        {
            for (uint16 i=0; i < 8; ++i)
            {
                build_sh_octree(octree, target, center + extents * factors[i], extents*0.5f, camera, query, octree->m_nodes[cur_node].get_child()+i, depth+1);
            }
        }

    }
}

volatile bool update_octree = false;

static const float g_octree_size = 4.0f;


void gi_sample::Render()
{
    PROFILE_FUNC();

    if (RBInput::GetSingleton()->IsPressed(RBInput::IK_P))
    {
        create_gi();
    }

    RBApplication::GetSingleton()->OnIdle();

    RBApplication* app = RBApplication::GetSingleton();
    RBCamera* camera = &RBCamera::Get(app->get_camera());


    RBApplication::render_scene(app->get_scene(), camera, RBRenderer::GetSingleton(), RBRenderer::GetSingleton()->GetDefaultGroup());
    static float3 cam_pos(camera->WorldXForm().Translate);
    if (update_octree)
    {
        cam_pos = camera->WorldXForm().Translate;
        update_octree = false;
    }

    draw_aab_count = 0;
    {
        PROFILE_BLOCK(draw_octree);
        draw_octree(&m_private_members->octree, cam_pos, float3(g_octree_size,g_octree_size,g_octree_size));
    }
    RBRenderPassDebug::Render(RBRenderer::GetSingleton());

    //render_gi(camera);

    if (dump_profiler)
    {
        RBApplication::GetSingleton()->dump_profiler();
        dump_profiler = false;
    }


    PROFILE_CLEAR();

    RBRenderer::GetSingleton()->Swap();
}

gi_sample::gi_sample( void* hWnd, int width, int height )
: WIDTH(width)
, HEIGHT(height)
{
    m_private_members = new rb_private_members;
    ID3D11Device* dev;
    IDXGISwapChain* swap_chain;
    RBD3D11Renderer::create_device_swap_chain(&dev, &swap_chain, RBD3D11RendererInit(hWnd, width, height));
    new RBD3D11Renderer();
    RBApplication::GetSingleton()->Width = WIDTH;
    RBApplication::GetSingleton()->Height = HEIGHT;
    RBD3D11Renderer::GetSingleton()->init(dev, swap_chain);
    RBApplication::GetSingleton()->Init(RBD3D11Renderer::GetSingleton());

    init_gi();

    D3D11_QUERY_DESC query_desc;
    query_desc.MiscFlags = 0;
    query_desc.Query = D3D11_QUERY_OCCLUSION;
    RBD3D11Renderer::GetSingleton()->GetDevice()->CreateQuery(&query_desc, &m_private_members->m_query);

}

void gi_sample::Release()
{
    m_private_members->m_query->Release();
    RBApplication::GetSingleton()->dump_profiler();

    RBApplication::GetSingleton()->Release();
    delete RBApplication::GetSingleton();
}


void gi_sample::AddFile( const char* filename )
{
    RBApplication::GetSingleton()->Open(filename);
}

void gi_sample::ChangeSize( int width, int height )
{
    WIDTH = width;
    HEIGHT = height;

    RBD3D11Renderer::GetSingleton()->change_size(width, height);

    RBApplication::GetSingleton()->Width = WIDTH;
    RBApplication::GetSingleton()->Height = HEIGHT;


}

void gi_sample::create_gi()
{
    RBHandle<RBCamera> camera = RBApplication::GetSingleton()->get_camera();
    RBCamera new_camera(InvalidHandle);
    {
        camera->copy_members(&new_camera);
        new_camera.SetViewFrustum(-1.0f, 1.0f, 1.0f, -1.0f, 0.1f, 10.0f, false);
    }

    ID3D11Query* query;
    D3D11_QUERY_DESC query_desc;
    query_desc.Query = D3D11_QUERY_OCCLUSION;
    query_desc.MiscFlags = 0;
    RBD3D11Renderer::GetSingleton()->GetDevice()->CreateQuery(&query_desc, &query);
    m_private_members->octree.clear();
    build_sh_octree(&m_private_members->octree, m_gi_target, camera->WorldXForm().Translate, g_octree_size, &new_camera, query);
    update_octree = true;
    query->Release();

#if 0
    static const int order = 3;
    float  sh_r[order*order] = {0};
    float  sh_g[order*order] = {0};
    float  sh_b[order*order] = {0};

    static const RBColorA clr_colors[] = 
    {
        RBColorA(1,0,0,1), RBColorA(0,0,0,1),
        RBColorA(0,1,0,1), RBColorA(1,1,0,1),
        RBColorA(0,0,1,1), RBColorA(1,1,1,1),
    };

    // Hide all the spheres before shooting cube maps
    for (uint32 ui=0; ui < m_private_members->m_debug_spheres.size(); ++ui)
    {
        m_private_members->m_debug_spheres[ui]->IsVisible = false;
    }

    // Shoot all the cube maps, and assign the SH coeffs to each sphere
    for (uint32 ui=0; ui < m_private_members->m_debug_spheres.size(); ++ui)
    {
        RBGeometryHandle dest_geom = m_private_members->m_debug_spheres[ui];

        float3 cam_pos = camera->WorldXForm().Translate + dest_geom->LocalXForm().Translate;
        dest_geom->WorldXForm().Translate = cam_pos;
#if 1
        for (uint32 ui=0; ui < 6; ++ui)
        {
            new_camera.LookAt(g_dirs[ui], g_ups[ui]);
            // Go left handed because we handle with a reflection
            {
                float* col_x = new_camera.WorldXForm().Rotate.GetCol(RBCamera::CD_RIGHT);
                col_x[0] = -col_x[0];
                col_x[1] = -col_x[1];
                col_x[2] = -col_x[2];
            }

            render_gi(&new_camera);
            UINT dst_sub_res = D3D10CalcSubresource(0, ui, m_textureGI_cpu->GetMipCount());

            RBD3D11Renderer::GetSingleton()->GetDeviceContext()->CopySubresourceRegion((ID3D11Resource*)m_textureGI_cpu->GetNative(), dst_sub_res, 0, 0, 0, (ID3D11Resource*)m_gi_target_rgb8->GetRenderTargetView()->GetBuffer()->GetNative(), 0, NULL);
        }

        memset(sh_r, 0, sizeof(sh_r));
        memset(sh_g, 0, sizeof(sh_g));
        memset(sh_b, 0, sizeof(sh_b));

        SH_Encode(m_textureGI_cpu, sh_r, sh_g, sh_b, order);


        for (uint32 ui=0; ui < 6; ++ui)
        {
            SH_Decode(m_textureGI_sh_cube, m_textureGI_sh_cpu_output, m_textureGI_cpu, sh_r, sh_g, sh_b, order,false);
        }


        float4* sh_rgb;
        uint32 count;
        dest_geom->DataContainer.get_effect_data<float4>("sh_rgb", &sh_rgb, &count);

        for (uint32 ui=0; ui < count; ++ui)
        {
            sh_rgb[ui](0) = sh_r[ui];
            sh_rgb[ui](1) = sh_g[ui];
            sh_rgb[ui](2) = sh_b[ui];
        }
#endif
    }

    RBBVolume::UpdateEntity(m_private_members->m_debug_sphere_entity);

    // Restore all sphere states
    for (uint32 ui=0; ui < m_private_members->m_debug_spheres.size(); ++ui)
    {
        m_private_members->m_debug_spheres[ui]->IsVisible = true;
    }
#endif
}

void gi_sample::render_gi(RBCamera* camera)
{
    RBRenderer* renderer = RBRenderer::GetSingleton();
    //RBApplication::render_scene(RBApplication::GetSingleton()->get_scene(), camera, RBD3D11Renderer::GetSingleton(), m_gi_target);
    RBCamera new_camera(InvalidHandle);
    {
        camera->copy_members(&new_camera);
        new_camera.SetViewFrustum(-1.0f, 1.0f, 1.0f, -1.0f, 0.1f, 10.0f, false);
    }

    render_cube(camera->WorldXForm().Translate, &new_camera, m_gi_target, m_private_members->m_query, 1<<0);
    m_fullscreen->Render(renderer, m_gi_target_rgb8);


    UINT64 queryData; // This data type is different depending on the query type
    while( S_OK != RBD3D11Renderer::GetSingleton()->GetDeviceContext()->GetData(m_private_members->m_query, &queryData, sizeof(UINT64), 0) )
    {
        ::Sleep(1);
    }

    pixels_visible = (int) queryData;

}

static IDirect3DDevice9Ex* g_dev_d3d = NULL;

void gi_sample::init_gi()
{
    static const uint32 g_res_gi = 128;
    m_gi_target = RBRenderTargetGroup::Create(g_res_gi, g_res_gi, PF_B8G8R8A8_UNORM, PF_D24_UNORM_S8_UINT);
    m_gi_target_rgb8 = RBRenderTargetGroup::Create(g_res_gi, g_res_gi, PF_R8G8B8A8_UNORM, PF_UNKNOWN);
    m_textureGI = (ID3D11Texture2D*)m_gi_target->GetRenderTargetView()->GetBuffer()->GetNative();
    m_fullscreen = new RBRenderPassFullScreen("Blit");
    RBShaderResourceView* shader_view = RBRenderer::GetSingleton()->GetHBMgr()->CreateShaderResourceView(m_gi_target->GetRenderTargetView()->GetBuffer());
    m_fullscreen->GetMaterial()->DataContainer.AddView("diffuse_texture", shader_view);
    m_fullscreen->GetMaterial()->DataContainer.GetSampler("diffuse_sampler");

    m_textureGI_cpu = RBRenderer::GetSingleton()->GetHBMgr()->CreateTexture(rb_hw_buffer_usage::HBU_CPUREAD | rb_hw_buffer_usage::HBU_GPUREAD | rb_hw_buffer_usage::HBU_GPUWRITE,
        g_res_gi, g_res_gi, RBPixelFormat::PF_R8G8B8A8_UNORM, RBTextureType::CUBE, 1);

    m_textureGI_sh_cpu_output = RBRenderer::GetSingleton()->GetHBMgr()->CreateTexture(rb_hw_buffer_usage::HBU_CPUWRITE | rb_hw_buffer_usage::HBU_GPUREAD | rb_hw_buffer_usage::HBU_SHADER_RESOURCE, m_textureGI_cpu->GetWidth(), m_textureGI_cpu->GetWidth(), PF_R8G8B8A8_UNORM, RBTextureType::SURFACE, 1);
    m_textureGI_sh_cube = RBRenderer::GetSingleton()->GetHBMgr()->CreateTexture(rb_hw_buffer_usage::HBU_GPUREAD | rb_hw_buffer_usage::HBU_GPUWRITE | rb_hw_buffer_usage::HBU_SHADER_RESOURCE, m_textureGI_cpu->GetWidth(), m_textureGI_cpu->GetWidth(), PF_R8G8B8A8_UNORM, RBTextureType::CUBE, 1);


    static const int order = 6;
    float  sh_r[order*order] = {0};
    float  sh_g[order*order] = {0};
    float  sh_b[order*order] = {0};
    SH_Encode(m_textureGI_cpu, sh_r, sh_g, sh_b, order);
    SH_Decode(m_textureGI_sh_cube, m_textureGI_sh_cpu_output, m_textureGI_cpu, sh_r, sh_g, sh_b, order, true);

    RBStream* pStream = RBStream::Load("Data/Sphere.dae");
    RBEffect* sphere_effect_sh = RBEffectFactory::GetSingleton()->RetrieveEffect("sphere_debug_sh");

    m_private_members->m_debug_sphere_entity = pStream->GetObjectAt(0);
#if 0
    RBEntityHandle sphere_entity = m_private_members->m_debug_sphere_entity;
    RBAssert(m_private_members->m_debug_sphere_entity != InvalidHandle, "unexpected missing data");
    {
        m_private_members->m_debug_sphere_template = sphere_entity->GetComponent<RBGeometry>(0);
        RBBVolumeHandle bounding_sphere_template = sphere_entity->GetComponent<RBBVolume>(0);
        bounding_sphere_template->Update(bounding_sphere_template->WorldXForm());
        RBAssert(bounding_sphere_template->GetObject<RBGeometry>() == m_private_members->m_debug_sphere_template, "unexpected bounding volume");
        if (m_private_members->m_debug_sphere_template != InvalidHandle)
        {
            RBMaterialPtr new_mat = new RBMaterial();
            m_private_members->m_debug_sphere_template->GetMaterial()->CopyMembers(new_mat);
            float4* sh_rgb = m_private_members->m_debug_sphere_template->DataContainer.AddEffectData<float4>("sh_rgb", order*order);

            new_mat->DataContainer.AddView("diffuse_texture", RBD3D11Renderer::GetSingleton()->GetHBMgr()->CreateShaderResourceView(m_textureGI_sh_cube));
            new_mat->DataContainer.GetSampler("diffuse_sampler");
            new_mat->SetEffect(sphere_effect_sh);
            m_private_members->m_debug_sphere_template->SetMaterial(new_mat);

            // Create all instances based on this template

            static const int extents = 6;

            for (int z=-extents; z<=extents; ++z)
            {
                for (int y=-extents; y<=extents; ++y)
                {
                    for (int x=-extents; x<=extents; ++x)
                    {
                        RBGeometryHandle dest_geom = RBGeometry::Create(sphere_entity);
                        m_private_members->m_debug_sphere_template->CopyMembers(dest_geom);
                        dest_geom->LocalXForm().Translate = 3.0f*float3(x,y,z);
                        m_private_members->m_debug_sphere_entity->AddComponent(dest_geom);
                        m_private_members->m_debug_spheres.push_back(dest_geom);

                        RBBVolumeHandle bounding_volume = RBBVolume::Create(sphere_entity, dest_geom->get_xform_handle());
                        bounding_sphere_template->copy_members(bounding_volume);
                        bounding_volume->Update(bounding_volume->WorldXForm());
                        bounding_volume->SetObject(dest_geom);
                        m_private_members->m_debug_sphere_entity->AddComponent(bounding_volume);
                    }
                }
            }
        }
        RBApplication::GetSingleton()->OnEntityLoaded(sphere_entity);
    }
#endif
}
